from model.mocov2 import MoCoV2_Model
import torch

if __name__ == '__main__':
    # x_q = torch.randn(2, 3, 32, 32)
    # x_k = torch.randn(2, 3, 32, 32)

    # class args:
    #     dataset = "cifar10"
    #     backbone = "resnet18"
    #     batch_size = 128
    #     num_classes=10
    #     out_dim=128

    # model = MoCoV2_Model(args)
    # logit, label = model(x_q, x_k)
    # print(logit.shape, label.shape)

    backbone_sim_matrix=torch.randn(512,9,9)
    
    densecl_sim_index=backbone_sim_matrix.max(dim=2)[1]# [512,9]

    k_grid=torch.randn(512,128,9)
    print(densecl_sim_index.unsqueeze(1).shape)
    index=densecl_sim_index.unsqueeze(1).expand(-1, k_grid.size(1), -1)
    print(index.shape)
    indexed_k_grid = torch.gather(k_grid, 2, densecl_sim_index.unsqueeze(1).expand(-1, k_grid.size(1), -1))

    print(indexed_k_grid.shape) #[512,128,9]
    print(indexed_k_grid)
